#!/usr/bin/python
# -*- coding: UTF-8 -*-

import pandas as pd
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split

# 加载数据
path = '../datas/household_power_consumption_1000.txt'
data = pd.read_csv(filepath_or_buffer=path, sep=";")

print(data.iloc[:,2:4].head(5))
# print(data.info())

# 切分数据
x=data.iloc[:,2:4]
y=data.iloc[:,5]

# 切分测试数据
x_train,x_text,y_train,y_test = train_test_split(x, y, train_size=0.8, random_state=0)
# print(x_train.shape)

X = np.mat(x_train)
Y = np.mat(y_train).reshape((-1, 1))
print(Y.shape)

theta = (X.T * X).I * X.T * Y
print(theta)

y_hat = np.mat(x_text) * theta

plt.figure(facecolor='w')
t = np.arange(len(x_text))
plt.plot(t,y_test,'r-',label='真实值')
plt.plot(t,y_hat,'b-',label="预测值")
plt.legend('lower right')
plt.show()

if __name__ == '__main__':
    pass